import numpy as np
import torch
import transformers
from transformers import LlamaForCausalLM, LlamaTokenizer
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from safetensors import safe_open
import cv2
from safetensors.torch import save_file

par_dict={}
par_dict2={}
llamaLayersNum=32
llama_layer_name_suffix_dict = {
    "q": ".self_attn.q_proj.weight",
    "k": ".self_attn.k_proj.weight",
    "v": ".self_attn.v_proj.weight",
    "o": ".self_attn.o_proj.weight",
    'mlp_down': ".mlp.down_proj.weight",
    'mlp_gate': ".mlp.gate_proj.weight",
    'mlp_up': ".mlp.up_proj.weight",
}

llama_layer_name_suffix_list = [
    ".self_attn.q_proj.weight",
    ".self_attn.k_proj.weight",
    ".self_attn.v_proj.weight",
    ".self_attn.o_proj.weight",
    ".mlp.down_proj.weight",
    ".mlp.gate_proj.weight",
    ".mlp.up_proj.weight",
]


resList=[]#按照llama_layer_name_suffix_list中的顺序排列的权重矩阵





def store2ParDict(model_info):
    global par_dict
    for i in range(model_info['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info['PathPrefix'] + s + model_info['PathSuffix']
        print(path)
        getModel(path)

def getModel(model_path):
    global par_dict
    with safe_open(model_path, framework="pt", device='cpu') as f:
        for k in f.keys():
            par_dict[k] = f.get_tensor(k)


def showParDict():
    global par_dict
    for i, name in enumerate(par_dict):
        parameter = par_dict[name].numpy()
        print(name, parameter.shape)


def getLayerInfo(txtName,model_path):
    # 输出模型每一层的相关信息到txtName.txt文件中，model_path为模型文件的权重路径
    def eleNum(a):
        if a.ndim == 1:
            return a.shape[0]
        return a.shape[0] * a.shape[1]

    def func(a, avg, std):
        pd1 = avg + std  # 距离在正的一个方差之内 positive delta
        pd2 = avg + 2 * std
        pd3 = avg + 3 * std
        nd1 = avg - std
        nd2 = avg - 2 * std
        nd3 = avg - 3 * std
        count3 = np.sum((a >= nd3) * (a <= pd3))
        count2 = np.sum((a >= nd2) * (a <= pd2))
        count1 = np.sum((a >= nd1) * (a <= pd1))
        return count1, count2, count3

    par_dict=getModel(model_path)
    f=open(txtName+'weight.txt','a')
    for i,name in enumerate(par_dict):
        parameter=par_dict[name].numpy()
        avg = np.average(parameter)
        std = np.std(parameter)
        max = np.max(parameter)
        min = np.min(parameter)
        count1, count2, count3 = func(parameter, avg, std)
        elesum = eleNum(parameter)
        print(i)
        print(name, parameter.shape, file=f)
        print("avg:%f std:%f max:%f min:%f" % (avg, std, max, min), file=f)
        print("一个方差内的数据个数:%d 占比:%f" % (count1, float(count1 / elesum)), file=f)
        print("两个方差内的数据个数:%d 占比:%f" % (count2, float(count2 / elesum)), file=f)
        print("三个方差内的数据个数:%d 占比:%f" % (count3, float(count3 / elesum)), file=f)
        print("", file=f)


def getLayerName(txtName,model_path):
    # 输出模型每一层的name与size到txtName.txt文件中，model_path为模型文件的权重路径
    def eleNum(a):
        if a.ndim == 1:
            return a.shape[0]
        return a.shape[0] * a.shape[1]

    getModel(model_path)
    f=open(txtName+'name.txt','a')
    for i, name in enumerate(par_dict):
        parameter = par_dict[name].numpy()
        elesum = eleNum(parameter)
        print(i)
        print(name, parameter.shape, file=f)
        print("", file=f)


def getSingleImg(model_path):
    global res
    getModel(model_path)
    for i,name in enumerate(par_dict):
        parameter=par_dict[name].numpy()
    if parameter.ndim>=2 and ('model.layers.0.self_attn.o_proj.weight' in name):
        avg = np.average(parameter)
        std = np.std(parameter)
        #白色为1，黑色为0
        img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
        res=img
    elif parameter.ndim>=2 and ('self_attn.o_proj.weight' in name):
        print(i)
        print(res.shape)
        avg = np.average(parameter)
        std = np.std(parameter)
        img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
        res = np.concatenate([res, img], axis=1)
        # axis=1横向拼接
        # axis=0纵向拼接


def initResList():
    global  resList

    for index,suffix_name in enumerate(llama_layer_name_suffix_list):
        resList.append(np.zeros([1,1]))

def getBinaryImg(parameter):#1为白色，0为黑色
    #大于零为1，小于零为0
    # img=255-255*(parameter>=0)#原矩阵中大于零白色，小于零黑色
    # 大于零为1，小于零为0

    #一个方差内为0，其余为1
    avg = np.average(parameter)
    std = np.std(parameter)
    img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
    #一个方差内为0，其余为1

    # 两个个方差内为0，其余为1
    # avg = np.average(parameter)
    # std = np.std(parameter)
    # img = 255 - 255 * ((parameter >= avg - 2*std) * (parameter <= avg + 2*std))
    # 两个方差内为0，其余为1
    return img

def getSingleImg2(model_info):#这里传入的是两个矩阵的差值部分
    global resList,par_dict,llama_layer_name_suffix_list
    store2ParDict(model_info)
    showParDict()
    initResList()
    for layerCount in range(llamaLayersNum):
        for index,suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount<=0:#只有符合条件的index对应的矩阵种类会出现在resList结果中
                #只需要在这里修改index部分的代码即可
                name='model.layers.' + str(layerCount)+suffix_name
                print(layerCount)
                print(index)
                print(name)
                cur_matrix=par_dict[name].numpy()
                if layerCount==0:
                    resList[index]=getBinaryImg(cur_matrix)
                else:
                    resList[index]=np.concatenate([resList[index], getBinaryImg(cur_matrix)], axis=1)
                    print(resList[index].shape)
                    # axis=1横向拼接
                    # axis=0纵向拼接


def getImg(model_info):#getSingleImg2函数得到的结果
    global resList

    getSingleImg2(model_info)
    cv2.imwrite("self_attn.weight.o0.png", resList[3].astype(np.uint8))


def getBinary(parameter,avg,std):
    img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
    return img

def getDiffImg():
    global resList,par_dict,par_dict2,llama_layer_name_suffix_list
    showParDict()
    initResList()
    for layerCount in range(llamaLayersNum):
        for index,suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount<=0:#只有符合条件的index对应的矩阵种类会出现在resList结果中
                #只需要在这里修改index部分的代码即可
                name='model.layers.' + str(layerCount)+suffix_name
                print(layerCount)
                print(index)
                print(name)

                cur_matrix = par_dict[name].numpy() - par_dict2[name].numpy()
                # print(np.average(np.abs(np.divide(par_dict[name].numpy(),cur_matrix))))
                avg = np.average(cur_matrix)
                std = np.std(cur_matrix)
                # print(par_dict[name].numpy())
                # print(par_dict2[name].numpy())
                # print(cur_matrix)
                # print(getBinary(cur_matrix,avg,std))
                #原始llama-微调llama模型的差值
                if layerCount==0:
                    # resList[index]=getBinaryImg(cur_matrix)
                    resList[index] = getBinary(cur_matrix,avg,std)
                else:
                    resList[index]=np.concatenate([resList[index], getBinary(cur_matrix,avg,std)], axis=1)
                    print(resList[index].shape)
                    # axis=1横向拼接
                    # axis=0纵向拼接


def getDiffValue():
    def eleNum(a):
        if a.ndim == 1:
            return a.shape[0]
        return a.shape[0] * a.shape[1]

    def func(a, avg, std):
        pd1 = avg + std  # 距离在正的一个方差之内 positive delta
        pd2 = avg + 2 * std
        pd3 = avg + 3 * std
        nd1 = avg - std
        nd2 = avg - 2 * std
        nd3 = avg - 3 * std
        count3 = np.sum((a >= nd3) * (a <= pd3))
        count2 = np.sum((a >= nd2) * (a <= pd2))
        count1 = np.sum((a >= nd1) * (a <= pd1))
        return count1, count2, count3

    global resList,par_dict,par_dict2,llama_layer_name_suffix_list
    showParDict()
    f = open('llama7bdiffweight.txt', 'a')
    for layerCount in range(llamaLayersNum):
        for index,suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount<=2:#只有符合条件的index对应的矩阵种类会出现在resList结果中
                #只需要在这里修改index部分的代码即可
                name='model.layers.' + str(layerCount)+suffix_name
                print(layerCount)
                print(index)
                print(name)
                parameter = par_dict[name].numpy() - par_dict2[name].numpy()
                # print(np.average(np.abs(np.divide(par_dict[name].numpy(),cur_matrix))))
                # avg = np.average(parameter)
                avgabs = np.average(np.abs(parameter))
                # std = np.std(parameter)
                # max = np.max(parameter)
                # min = np.min(parameter)
                # count1, count2, count3 = func(parameter, avg, std)
                # elesum = eleNum(parameter)
                print(name, parameter.shape, file=f)
                # print("avg:%f std:%f max:%f min:%f" % (avgabs, std, max, min), file=f)
                print("avgabs:%f" % (avgabs), file=f)
                # print("一个方差内的数据个数:%d 占比:%f" % (count1, float(count1 / elesum)), file=f)
                # print("两个方差内的数据个数:%d 占比:%f" % (count2, float(count2 / elesum)), file=f)
                # print("三个方差内的数据个数:%d 占比:%f" % (count3, float(count3 / elesum)), file=f)
                print("", file=f)


def getDiff(model_info,model_info2):
    def eleNum(a):
        if a.ndim == 1:
            return a.shape[0]
        return a.shape[0] * a.shape[1]

    def func(a, avg, std):
        pd1 = avg + std  # 距离在正的一个方差之内 positive delta
        pd2 = avg + 2 * std
        pd3 = avg + 3 * std
        nd1 = avg - std
        nd2 = avg - 2 * std
        nd3 = avg - 3 * std
        count3 = np.sum((a >= nd3) * (a <= pd3))
        count2 = np.sum((a >= nd2) * (a <= pd2))
        count1 = np.sum((a >= nd1) * (a <= pd1))
        return count1, count2, count3

    global resList, par_dict, par_dict2, llama_layer_name_suffix_list

    for i in range(model_info2['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info2['PathPrefix'] + s + model_info2['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            for k in f.keys():
                par_dict2[k] = f.get_tensor(k)
                print(f.metadata())

    for i in range(model_info['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info['PathPrefix'] + s + model_info['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            for k in f.keys():
                par_dict[k] = f.get_tensor(k)
        #原始llama-微调llama模型的差值

        showParDict()
        f = open('llama7bdiffweight.txt', 'a')
        for layerCount in range(llamaLayersNum):
            for index, suffix_name in enumerate(llama_layer_name_suffix_list):

                if layerCount <= 2:  # 只有符合条件的index对应的矩阵种类会出现在resList结果中
                    # 只需要在这里修改index部分的代码即可
                    name = 'model.layers.' + str(layerCount) + suffix_name
                    print(layerCount)
                    print(index)
                    print(name)
                    parameter = par_dict[name].numpy() - par_dict2[name].numpy()
                    # print(np.average(np.abs(np.divide(par_dict[name].numpy(),cur_matrix))))
                    # avg = np.average(parameter)
                    avgabs = np.average(np.abs(parameter))
                    # std = np.std(parameter)
                    # max = np.max(parameter)
                    # min = np.min(parameter)
                    # count1, count2, count3 = func(parameter, avg, std)
                    # elesum = eleNum(parameter)
                    print(name, parameter.shape, file=f)
                    # print("avg:%f std:%f max:%f min:%f" % (avgabs, std, max, min), file=f)
                    print("avgabs:%f" % (avgabs), file=f)
                    # print("一个方差内的数据个数:%d 占比:%f" % (count1, float(count1 / elesum)), file=f)
                    # print("两个方差内的数据个数:%d 占比:%f" % (count2, float(count2 / elesum)), file=f)
                    # print("三个方差内的数据个数:%d 占比:%f" % (count3, float(count3 / elesum)), file=f)
                    print("", file=f)

    # getDiffValue()
    getDiffImg()
    for index,suffix_name in enumerate(llama_layer_name_suffix_list):
        name=str(index)+suffix_name+'_diff.png'
        cv2.imwrite(name,resList[index].astype(np.uint8))


# def outputDiffmatrix(diff):
#     np.set_printoptions(threshold=np.inf)
#     f = open('llama7bdiffmatrix.txt','a')

def getDiffMat(model_info,model_info2):#得到插值矩阵，并按照条带状分布来对原始llama矩阵进行相加，比较说明得到的模型可以输出中文
    def getOneZero(parameter):
        avg = np.average(parameter)
        std = np.std(parameter)
        # img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
        img = 1-1 * ((parameter >= avg - std) * (parameter <= avg + std))#结果为1的地方表示差距大，结果为0的地方表示差距小
        return img

    global resList, par_dict, par_dict2
    medata={}
    for i in range(model_info2['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info2['PathPrefix'] + s + model_info2['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())
            medata=f.metadata()
            print(medata)
            for k in f.keys():
                par_dict2[k] = f.get_tensor(k)

    for i in range(model_info['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info['PathPrefix'] + s + model_info['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())
            for k in f.keys():
                par_dict[k] = f.get_tensor(k)
        # 原始llama-微调llama模型的差值
    print(2)

    k=20
    for layerCount in range(llamaLayersNum):
        for index, suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount <=31:  # 只有符合条件的index对应的矩阵种类会出现在resList结果中
                # 只需要在这里修改index部分的代码即可
                name = 'model.layers.' + str(layerCount) + suffix_name
                print(name)
                diff = par_dict[name].numpy() - par_dict2[name].numpy()#原始llama减去微调llama
                signMitrix=1*(par_dict[name].numpy()>par_dict2[name].numpy())#原始llama大于微调llama的部分为1
                signMitrix[signMitrix==0]=-1#如果差值为负，那么signMirix对应元素为-1否则为1
                avgabs = np.average(np.abs(diff))
                oneZeroDiff=getOneZero(diff)
                mask=avgabs*oneZeroDiff*signMitrix
                matrix=par_dict[name].numpy()
                matrix=matrix+k*mask#这里用加法才对
                par_dict[name]=torch.from_numpy(matrix)
                par_dict[name] = par_dict[name].half()
    print("begin save")
    print(medata)
    save_file(par_dict, './04.safetensors',medata)
    print("finish save")



def getBigDiffImg(model_info,model_info2):
    def getOneZero(parameter):
        avg = np.average(parameter)
        std = np.std(parameter)
        # img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
        img = 1-1 * ((parameter >= avg - std) * (parameter <= avg + std))#结果为1的地方表示差距大，结果为0的地方表示差距小
        return img
    def getBinaryByOriginStd(parameter,avg,std):
        return 255*((parameter >= avg - std) * (parameter <= avg + std))
        # 让黑色表示更重要的值，即一个方差内的值表示白色，也即一个方差内值为255

    global resList, par_dict, par_dict2
    for i in range(model_info2['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info2['PathPrefix'] + s + model_info2['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())

            for k in f.keys():
                par_dict2[k] = f.get_tensor(k)

    for i in range(model_info['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info['PathPrefix'] + s + model_info['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())
            for k in f.keys():
                par_dict[k] = f.get_tensor(k)
        # 原始llama-微调llama模型的差值
    print(2)

    showParDict()
    initResList()
    k=20
    for layerCount in range(llamaLayersNum):
        for index, suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount <=4:  # 只有符合条件的index对应的矩阵种类会出现在resList结果中
                # 只需要在这里修改index部分的代码即可
                name = 'model.layers.' + str(layerCount) + suffix_name
                print(name)
                diff = par_dict[name].numpy() - par_dict2[name].numpy()#原始llama减去微调llama

                avg=np.average(par_dict[name].numpy())
                std=np.std(par_dict[name].numpy())

                signMitrix = 1 * (par_dict[name].numpy() > par_dict2[name].numpy())  # 原始llama大于微调llama的部分为1
                signMitrix[signMitrix == 0] = -1  # 如果差值为负，那么signMirix对应元素为-1否则为1

                avgabs = np.average(np.abs(diff))
                oneZeroDiff = getOneZero(diff)#权重脑图掩码
                mask = avgabs * oneZeroDiff * signMitrix

                cur_matrix = par_dict[name].numpy()
                cur_matrix = cur_matrix + k*mask  # 这里用加法才对


                if layerCount==0:
                    # resList[index]=getBinaryImg(cur_matrix)
                    # resList[index] = getBinary(cur_matrix,avg,std)
                    resList[index] = getBinaryByOriginStd(cur_matrix, avg, std)
                else:
                    resList[index]=np.concatenate([resList[index], getBinaryByOriginStd(cur_matrix, avg, std)], axis=1)
                    print(resList[index].shape)

    for index,suffix_name in enumerate(llama_layer_name_suffix_list):
        name=str(index)+suffix_name+'_diff.png'
        cv2.imwrite(name,resList[index].astype(np.uint8))


def getWeightMapDiff(model_info,model_info2):
    def get255Binary(parameter):
        avg = np.average(parameter)
        std = np.std(parameter)
        # img = 255 - 255 * ((parameter >= avg - std) * (parameter <= avg + std))
        img = 255* ((parameter >= avg - std) * (parameter <= avg + std))#结果为黑色的部分就是要调整的部分
        return img

    global resList, par_dict, par_dict2
    for i in range(model_info2['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info2['PathPrefix'] + s + model_info2['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())

            for k in f.keys():
                par_dict2[k] = f.get_tensor(k)

    for i in range(model_info['FileNum']):
        if i < 9:
            s = '0' + str(i + 1)
        else:
            s = str(i + 1)
        path = model_info['PathPrefix'] + s + model_info['PathSuffix']
        print(path)
        with safe_open(path, framework="pt", device='cpu') as f:
            # print(f.metadata())
            for k in f.keys():
                par_dict[k] = f.get_tensor(k)
        # 原始llama-微调llama模型的差值
    print(2)

    showParDict()
    initResList()

    for layerCount in range(llamaLayersNum):
        for index, suffix_name in enumerate(llama_layer_name_suffix_list):

            if layerCount <= 0:  # 只有符合条件的index对应的矩阵种类会出现在resList结果中
                # 只需要在这里修改index部分的代码即可
                name = 'model.layers.' + str(layerCount) + suffix_name
                print(name)
                diff = par_dict[name].numpy() - par_dict2[name].numpy()  # 原始llama减去微调llama
                oneZeroDiff = get255Binary(diff)  # 权重脑图掩码


                # avg = np.average(cur_matrix)
                # std = np.std(cur_matrix)
                if layerCount == 0:
                    # resList[index]=getBinaryImg(cur_matrix)
                    resList[index] = oneZeroDiff
                else:
                    resList[index] = np.concatenate([resList[index], oneZeroDiff], axis=1)
                    print(resList[index].shape)

    for index, suffix_name in enumerate(llama_layer_name_suffix_list):
        name = str(index) + suffix_name + '_diff.png'
        cv2.imwrite(name, resList[index].astype(np.uint8))


if __name__=='__main__':
    originLlamaModelInfo = {}
    originLlamaModelInfo['PathPrefix']=r"/gdata/hpdesktop/webui/models/Llama-2-7b-hf/model-000"
    originLlamaModelInfo['PathSuffix']="-of-00002.safetensors"
    originLlamaModelInfo['FileNum']=2
    originLlamaModelInfo['ModelName']='llama'

    chatLlamaModelInfo = {}
    chatLlamaModelInfo['PathPrefix'] = r"/gdata/gpu/Llama-2-7b-chat-hf/model-000"
    chatLlamaModelInfo['PathSuffix'] = "-of-00002.safetensors"
    chatLlamaModelInfo['FileNum'] = 2
    chatLlamaModelInfo['ModelName'] = 'llamachat'


    loraModelInfo={}
    loraModelInfo['PathPrefix'] = r"/gdata/hpdesktop/webui/models/Llama-2-7b-hf-OA2000/model-000"
    loraModelInfo['PathSuffix'] = "-of-00006.safetensors"
    loraModelInfo['FileNum'] = 6
    loraModelInfo['ModelName'] = 'lora'

    mixtralInstructInfo={}
    mixtralInstructInfo['PathPrefix'] = r"/gdata/gpu/textgen3/models/Mixtral-7B_8/model-000"
    mixtralInstructInfo['PathSuffix'] = "-of-00019.safetensors"
    mixtralInstructInfo['FileNum'] = 19
    mixtralInstructInfo['ModelName'] = 'mixtralinstruct'

    changeModelInfo = {}
    changeModelInfo['PathPrefix'] = r"/gdata/gpu8613/changeWeightTest/"
    changeModelInfo['PathSuffix'] = ".safetensors"
    changeModelInfo['FileNum'] = 1
    changeModelInfo['ModelName'] = 'changeModel'

    lora10000ModelInfo = {}
    lora10000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/webui/models/Llama-2-7b-hf-OAC10000/model-000"
    lora10000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora10000ModelInfo['FileNum'] = 6
    lora10000ModelInfo['ModelName'] = 'lora'

    lora1000ModelInfo = {}
    lora1000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/1000/model-000"
    lora1000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora1000ModelInfo['FileNum'] = 6
    lora1000ModelInfo['ModelName'] = 'lora'

    lora2000ModelInfo = {}
    lora2000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/2000/model-000"
    lora2000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora2000ModelInfo['FileNum'] = 6
    lora2000ModelInfo['ModelName'] = 'lora'

    lora3000ModelInfo = {}
    lora3000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/3000/model-000"
    lora3000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora3000ModelInfo['FileNum'] = 6
    lora3000ModelInfo['ModelName'] = 'lora'

    lora4000ModelInfo = {}
    lora4000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/4000/model-000"
    lora4000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora4000ModelInfo['FileNum'] = 6
    lora4000ModelInfo['ModelName'] = 'lora'

    lora5000ModelInfo = {}
    lora5000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/5000/model-000"
    lora5000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora5000ModelInfo['FileNum'] = 6
    lora5000ModelInfo['ModelName'] = 'lora'

    lora6000ModelInfo = {}
    lora6000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel/6000/model-000"
    lora6000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora6000ModelInfo['FileNum'] = 6
    lora6000ModelInfo['ModelName'] = 'lora'

    lora2_1000ModelInfo = {}
    lora2_1000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel2/1000/model-000"
    lora2_1000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora2_1000ModelInfo['FileNum'] = 6
    lora2_1000ModelInfo['ModelName'] = 'lora'

    lora2_2000ModelInfo = {}
    lora2_2000ModelInfo['PathPrefix'] = r"/gdata/hpdesktop/zym/LoraModel2/2000/model-000"
    lora2_2000ModelInfo['PathSuffix'] = "-of-00006.safetensors"
    lora2_2000ModelInfo['FileNum'] = 6
    lora2_2000ModelInfo['ModelName'] = 'lora'
    # doSomething(model_info=originLlamaModelInfo)
    # getImg(model_info=originLlamaModelInfo)
    # getDiff(originLlamaModelInfo,chatLlamaModelInfo)
    # getDiffMat(lora1000ModelInfo,lora2000ModelInfo)
    # getBigDiffImg(originLlamaModelInfo,lora1000ModelInfo)
    # getBigDiffImg(originLlamaModelInfo,changeModelInfo)
    getWeightMapDiff(lora2_1000ModelInfo,lora2_2000ModelInfo)
